clc
clear

nVar  = 5;
nlev  = 100;
stencil_width = 5;
recBdy = ( stencil_width - 1 ) / 2;

q      = sym('q'     ,[nVar nlev]);
sqrtGL = sym('sqrtGL',[1 nlev]);
sqrtGR = sym('sqrtGR',[1 nlev]);
G13L   = sym('G13L'  ,[1 nlev]);
G13R   = sym('G13R'  ,[1 nlev]);

for iVar = 1:nVar
    disp(['Reconstructing variable No.',num2str(iVar)]);
    parfor iCell = 1:nlev
        q_weno = sym(nan(1,stencil_width));
        if iCell<=recBdy
            q_weno(1:recBdy-iCell+1) = nan;
            q_weno(recBdy-iCell+2:stencil_width) = q(iVar,1:iCell+recBdy);
        elseif iCell>=nlev-recBdy+1
            q_weno(nlev-iCell+stencil_width-recBdy+1:stencil_width) = nan;
            q_weno(1:nlev-iCell+stencil_width-recBdy) = q(iVar,iCell-recBdy:nlev);
        else
            q_weno = q(iVar,iCell-recBdy:iCell+recBdy);
        end
        
        qL(iVar,iCell) = WENO(q_weno,-1);
        qR(iVar,iCell) = WENO(q_weno, 1);
    end
end

F = sym(zeros(nVar,nlev+1));
parfor iEdge = 1:nlev
    disp(['Calculating flux on edge ',num2str(iEdge)]);
    if iEdge==1
        qL_edge     = qL(:,iEdge);
        sqrtGL_edge = sqrtGL(iEdge);
        G13L_edge   = G13L(iEdge);
    else
        qL_edge     = qR(:,iEdge-1);
        sqrtGL_edge = sqrtGR(iEdge-1);
        G13L_edge   = G13R(iEdge-1);
    end
    
    if iEdge==nlev+1
        qR_edge     = qR(:,nlev);
        sqrtGR_edge = sqrtGR(nlev);
        G13R_edge   = G13R(nlev);
    else
        qR_edge     = qL(:,iEdge);
        sqrtGR_edge = sqrtGL(iEdge);
        G13R_edge   = G13L(iEdge);
    end
    
    F(:,iEdge) = Flux(qL_edge,qR_edge,sqrtGL_edge,sqrtGR_edge,G13L_edge,G13R_edge);
end
iEdge = nlev+1;
disp(['Calculating flux on edge ',num2str(iEdge)]);
qL_edge     = qR(:,iEdge-1);
sqrtGL_edge = sqrtGR(iEdge-1);
G13L_edge   = G13R(iEdge-1);

qR_edge     = qR(:,nlev);
sqrtGR_edge = sqrtGR(nlev);
G13R_edge   = G13R(nlev);

F(:,iEdge) = Flux(qL_edge,qR_edge,sqrtGL_edge,sqrtGR_edge,G13L_edge,G13R_edge);

% F = Flux(qL,qR,sqrtGL,sqrtGR,G13L,G13R);

function F = Flux(qL,qR,sqrtGL,sqrtGR,G13L,G13R)

Kp    = 1/4;
Ku    = 3/4;
sigma = 1;

sp = 1;
sn = -1;

rhoL = density(qL);
rhoR = density(qR);
uL = vertical_speed(qL,sqrtGL,G13L) * sqrtGL;
uR = vertical_speed(qR,sqrtGR,G13R) * sqrtGR;

cL = acoustic_speed(qL,sqrtGL,G13L) * sqrtGL;
cR = acoustic_speed(qR,sqrtGR,G13R) * sqrtGR;

rhoh = ( rhoL + rhoR ) / 2;
a    = ( cL + cR ) / 2;

ML = uL / a;
MR = uR / a;

Mbar2 = ( uL^2 + uR^2 ) / ( 2 * a^2 );

pL = pressure(qL,sqrtGL);
pR = pressure(qR,sqrtGR);

Mh1 = M4(ML,sp) + M4(MR,sn) - Kp * ( 1 - sigma * Mbar2 ) * ( pR - pL ) / ( rhoh * a^2 );
Mh2 = M4(ML,sp) + M4(MR,sn);
Mh  = upwind( Mh1, Mh2, 1-sigma*Mbar2 );

coefL = sqrtGL * G13L;
coefR = sqrtGR * G13R;

ph = P5(ML,sp) * coefL * pL + P5(MR,sn) * coefR * pR - Ku * P5(ML,sp) * P5(MR,sn) * ( rhoL * coefL + rhoR * coefR ) * a * ( uR - uL );

p = P5(ML,sp) * pL + P5(MR,sn) * pR - Ku * P5(ML,sp) * P5(MR,sn) * ( rhoL + rhoR ) * a * ( uR - uL );

m = 0.5 * ( cL / sqrtGL + cR / sqrtGR ) * Mh;

F = 0.5 * m * ( qL + qR - sign(m) * ( qR - qL ) );
F(2) = F(2) + ph;
F(3) = F(3) + p;
end

function c = acoustic_speed(q,sqrtG,G13)
syms p0 Rd cpd cpv cvd cvv eqM Rd
% eqM = Md / Mv - 1 = 0.608
w1 = q(1);
w2 = q(2);
w3 = q(3);
w4 = q(4);
w5 = q(5);

coef1 = sqrt( p0*sqrtG*(1 + G13^2*sqrtG^2)*w1^2*w4^2*(w1 + w5)^3*(cpd*w1 + cpv*w5)*(cvd*w1 + cvv*w5)^3*  ...
    (w1 + w5 + eqM*w5)^2*((Rd*w4*(w1 + w5 + eqM*w5))/(p0*sqrtG*w1))^ ...
    ((cpd*w1 + cpv*w5)/(cvd*w1 + cvv*w5)) );

coef2 = sqrtG*w1*w4*(w1 + w5)^2*(cvd*w1 + cvv*w5)^2*(w1 + w5 + eqM*w5);

c = coef1 / coef2;
end

function rho = density(q)
rho = q(1) + q(5);
end

function p = pressure(q,sqrtG)
syms p0 Rd cpd cpv cvd cvv eqM Rd

w1 = q(1);
w2 = q(2);
w3 = q(3);
w4 = q(4);
w5 = q(5);

r = w5 / w1;
sh = r / ( 1 + r );
R = ( 1 + eqM * sh) * Rd;
cp = cpd + (cpv - cpd)*sh;
cv = cvd + (cvv - cvd)*sh;
kappa = cp / cv;
sqrtGrho = w1 + w5;
theta = w4 / w1;
rhotheta = sqrtGrho * theta / sqrtG;
p = p0 * ( rhotheta * R / p0 )^kappa;
end

function eta = vertical_speed(q,sqrtG,G13)
w1 = q(1);
w2 = q(2);
w3 = q(3);
w4 = q(4);
w5 = q(5);

u = w2 / ( w1 + w5 );
w = w3 / ( w1 + w5 );

eta = ( w + sqrtG * G13 * u ) / sqrtG;

end

function m2 = M2(M,s)
m2 = s / 4 * ( M + s )^2;
end

function m4 = M4(M,s)
beta = 1/8;
case1 = 0.5 * ( M + s )^2;
case2 = M2(M,s) * ( 1 - s * 16 * beta * M2(M,-s) );

criteria = abs(M) - 1;

m4 = upwind(case1,case2,criteria);

end

function p5 = P5(M,s)
alpha = 3/16;

case1 = 0.5 * ( 1 + s * sign(M) );
case2 = M2(M,s) * ( 2 * s - M ) - s * 16 * alpha * M * M2(M,-s);

criteria = abs(M) - 1;

p5 = upwind(case1,case2,criteria);
end

function q = upwind(qL,qR,criteria)
q = 0.5 * ( qL + qR - sign(criteria) * ( qR - qL ) );
end

function qrec = WENO(q,dir)
syms epsilon
C = sym('C',[1 3]);

qim(2) = q(1);
qim(1) = q(2);
qi     = q(3);
qip(1) = q(4);
qip(2) = q(5);

if(dir>0)
    stencil (1) =  qim(2)/3. - 7./6. * qim(1) + 11./6. * qi;
    stencil (2) = -qim(1)/6. + 5./6. * qi     +  1./3. * qip(1);
    stencil (3) =  qi    /3. + 5./6. * qip(1) -  1./6. * qip(2);
    
    coefA(1) = qim(2) - 2. * qim(1) + qi;
    coefA(2) = qim(1) - 2. * qi     + qip(1);
    coefA(3) = qi     - 2. * qip(1) + qip(2);
    
    coefB(1) =      qim(2) - 4. * qim(1) + 3. * qi;
    coefB(2) =      qim(1) -      qip(1);
    coefB(3) = 3. * qi     - 4. * qip(1) +      qip(2);
elseif(dir<0)
    stencil (1) =  qip(2)/3. - 7./6. * qip(1) + 11./6. * qi;
    stencil (2) = -qip(1)/6. + 5./6. * qi     +  1./3. * qim(1);
    stencil (3) =  qi    /3. + 5./6. * qim(1) -  1./6. * qim(2);
    
    coefA(1) = qip(2) - 2. * qip(1) + qi;
    coefA(2) = qip(1) - 2. * qi     + qim(1);
    coefA(3) = qi     - 2. * qim(1) + qim(2);
    
    coefB(1) =      qip(2) - 4. * qip(1) + 3. * qi;
    coefB(2) =      qip(1) -      qim(1);
    coefB(3) = 3. * qi     - 4. * qim(1) +      qim(2);
end

beta = coefA.^2 * 13. / 12. + coefB.^2 * 0.25;

if ~any(isnan(q))
    % WENO-Z
    tau  = abs( beta(3) - beta(1) );
    
    alpha = C .* ( 1 + tau ./ ( beta + epsilon ) );
else
    % WENO-JS
    alpha = C ./ ( beta + epsilon ).^2;
    for i = 1:3
        if dir>0
            if any(isnan(q(i:i+2)))
                alpha  (i) = 0;
                stencil(i) = 0;
            end
        elseif dir<0
            if any(isnan(q(i:i+2)))
                alpha  (4-i) = 0;
                stencil(4-i) = 0;
            end
        end
    end
    
end

omega = alpha ./ sum(alpha);

qrec = sum(stencil.*omega);
% qrec = dot(stencil, omega);

end
